import socket
import threading
import multiprocessing
import os
import selectors
from threading import Semaphore

# 全局变量,用于控制最大连接数
MAX_CONNECTIONS = 1000
connection_semaphore = Semaphore(MAX_CONNECTIONS)

def handle_client(key, mask, selector):
    """处理单个客户端连接"""
    sock = key.fileobj
    data = key.data
    if mask & selectors.EVENT_READ:
        try:
            recv_data = sock.recv(1024)
            if recv_data:
                response = f"服务器收到: {recv_data.decode()} (进程ID: {os.getpid()}, 线程ID: {threading.get_ident()})"
                data.outb += response.encode()
            else:
                print(f"关闭连接 {data.addr}")
                selector.unregister(sock)
                sock.close()
                connection_semaphore.release()
        except Exception as e:
            print(f"处理客户端时发生错误: {e}")
            selector.unregister(sock)
            sock.close()
            connection_semaphore.release()
    if mask & selectors.EVENT_WRITE:
        if data.outb:
            sent = sock.send(data.outb)
            data.outb = data.outb[sent:]

def accept_wrapper(sock, selector):
    conn, addr = sock.accept()
    print(f"新连接来自 {addr} (进程ID: {os.getpid()})")
    conn.setblocking(False)
    data = type('', (), {'addr': addr, 'inb': b'', 'outb': b''})()
    events = selectors.EVENT_READ | selectors.EVENT_WRITE
    selector.register(conn, events, data=data)

def process_worker(server_socket):
    """进程工作函数"""
    selector = selectors.DefaultSelector()
    selector.register(server_socket, selectors.EVENT_READ, data=None)

    while True:
        events = selector.select(timeout=None)
        for key, mask in events:
            if key.data is None:
                if connection_semaphore.acquire(blocking=False):
                    accept_wrapper(key.fileobj, selector)
                else:
                    print(f"拒绝新连接: 已达到最大连接数")
                    conn, addr = key.fileobj.accept()
                    conn.send("服务器繁忙,请稍后再试".encode())
                    conn.close()
            else:
                handle_client(key, mask, selector)

def start_server(host='127.0.0.1', port=8888, num_processes=4):
    server_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    server_socket.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
    server_socket.bind((host, port))
    server_socket.listen(5)
    server_socket.setblocking(False)
    print(f"服务器正在监听 {host}:{port}")
    print(f"最大连接数: {MAX_CONNECTIONS}")

    processes = []
    for _ in range(num_processes):
        process = multiprocessing.Process(target=process_worker, args=(server_socket,))
        process.start()
        processes.append(process)

    for process in processes:
        process.join()

if __name__ == "__main__":
    start_server()
